# 将数据集按8:1:1划分为3个集合

import os
import random
import shutil

# 源文件夹路径
old_folder = 'D:/上过的课/2024-2025-1-大四上/毕设/datasets/wash1+augmented/'
image_folder = old_folder+'images'
label_folder = old_folder+'labels'
ntlabel_folder = old_folder+ 'ntlabels'

# 目标文件夹路径
train_image_folder = old_folder+'train/images'
train_label_folder = old_folder+'train/labels'
train_ntlabel_folder = old_folder+'train/ntlabels'
val_image_folder = old_folder+'val/images'
val_label_folder = old_folder+'val/labels'
val_ntlabel_folder = old_folder+'val/ntlabels'
test_image_folder = old_folder+'test/images'
test_label_folder = old_folder+'test/labels'
test_ntlabel_folder = old_folder+'test/ntlabels'

# 创建目标文件夹
for folder in [train_image_folder, train_label_folder,train_ntlabel_folder, val_image_folder, val_label_folder, val_ntlabel_folder, test_image_folder, test_label_folder,test_ntlabel_folder]:
    if not os.path.exists(folder):
        os.makedirs(folder)

# 获取所有图片文件名
image_files = [f for f in os.listdir(image_folder) if f.endswith('.jpg')]
random.shuffle(image_files)

# 计算划分数量
total_count = len(image_files)
train_count = int(total_count * 0.8)
val_count = int(total_count * 0.1)
test_count = total_count - train_count - val_count

# 划分数据集
train_files = image_files[:train_count]
val_files = image_files[train_count:train_count + val_count]
test_files = image_files[train_count + val_count:]

# 定义复制函数
def copy_files(file_list, image_dest, label_dest, ntlabel_dest):
    for image_file in file_list:
        # 复制图片
        image_src_path = os.path.join(image_folder, image_file)
        image_dest_path = os.path.join(image_dest, image_file)
        shutil.copy2(image_src_path, image_dest_path)

        # 复制对应的标签文件
        label_file = os.path.splitext(image_file)[0] + '.txt'
        label_src_path = os.path.join(label_folder, label_file)
        label_dest_path = os.path.join(label_dest, label_file)
        if os.path.exists(label_src_path):
            shutil.copy2(label_src_path, label_dest_path)

        # 复制对应的no_two_wheeler标签文件
        label_file = os.path.splitext(image_file)[0] + '.txt'
        label_src_path = os.path.join(ntlabel_folder, label_file)
        label_dest_path = os.path.join(ntlabel_dest, label_file)
        if os.path.exists(label_src_path):
            shutil.copy2(label_src_path, label_dest_path)
            
# 复制文件到对应的文件夹
copy_files(train_files, train_image_folder, train_label_folder,train_ntlabel_folder)
copy_files(val_files, val_image_folder, val_label_folder,val_ntlabel_folder)
copy_files(test_files, test_image_folder, test_label_folder,test_ntlabel_folder)

print("数据集划分完成！")